class CliffWalkingEnv:
    """悬崖漫步环境"""

    def __init__(self, ncol=12, nrow=4):
        self.ncol = ncol  # 定义网格世界的列
        self.nrow = nrow  # 定义网格世界的行
        # 状态转移函数 P[state][action] = [(prob, next_state, reward)]
        self.P = [[[] for _ in range(4)] for _ in range(self.nrow * self.ncol)]

        self.setupEnv()

    def setupEnv(self):
        # 动作集合。change[0]:上,change[1]:下, change[2]:左, change[3]:右。坐标系原点(0,0)定义在左上角
        action = [[-1, 0], [1, 0], [0, -1], [0, 1]]
        for i in range(self.nrow):
            for j in range(self.ncol):
                for a in range(4):
                    # 当前已经处于悬崖和终点的情况，无法继续交互，任何动作的奖励都为 0
                    if i == self.nrow - 1 and j > 0:
                        self.P[self.state(i, j)][a] = [(1, self.state(i, j), 0, True)]
                        continue

                    # 获取下个状态
                    next_x = min(self.nrow - 1, max(0, i + action[a][0]))
                    next_y = min(self.ncol - 1, max(0, j + action[a][1]))
                    next_state = self.state(next_x, next_y)

                    # 只要下个状态没到终点或跌入悬崖，奖励默认为 -1
                    reward = -1
                    done = False

                    # 下个状态掉入悬崖或到达终点的情况
                    if next_x == self.nrow - 1 and next_y > 0:
                        done = True
                        if next_y == self.ncol - 1:  # 下个状态为终点
                            reward = 20
                        else:  # 下个状态为悬崖
                            reward = -100

                    self.P[self.state(i, j)][a] = [(1, next_state, reward, done)]

    def state(self, i, j):
        """将坐标转换为状态"""
        return i * self.ncol + j


def PrintPolicy(agent, action_meaning, disaster=[], end=[]):
    print("状态价值：")
    for i in range(agent.env.nrow):
        for j in range(agent.env.ncol):
            # 为了输出美观,保持输出6个字符
            print("%6.6s" % ("%.3f" % agent.v[i * agent.env.ncol + j]), end=" ")
        print()

    print("策略：")
    for i in range(agent.env.nrow):
        for j in range(agent.env.ncol):
            # 一些特殊的状态,例如悬崖漫步中的悬崖
            if (i * agent.env.ncol + j) in disaster:
                print("****", end=" ")
            elif (i * agent.env.ncol + j) in end:  # 目标状态
                print("EEEE", end=" ")
            else:
                a = agent.pi[i * agent.env.ncol + j]
                pi_str = ""
                for k in range(len(action_meaning)):
                    pi_str += action_meaning[k] if a[k] > 0 else "o"
                print(pi_str, end=" ")
        print()
